package org.idea.qiyu.framework.apm.agent.core.main;

import org.idea.qiyu.framework.apm.agent.core.dubbo.DubboAgent;
import org.idea.qiyu.framework.apm.agent.core.jdbc.DruidPoolAgent;
import org.idea.qiyu.framework.apm.agent.core.servlet.ServletAgent;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * ilink监控核型agent模块
 *
 * @Author linhao
 * @Date created in 7:05 上午 2021/9/6
 */
public class ILinkAgentCore {

    private static final String MOCK_CLUSTER_INVOKER = "org/apache/dubbo/rpc/cluster/support/wrapper/MockClusterInvoker";
    private static final String DRUID_POOL_CONN_CLASS = "com/alibaba/druid/pool/DruidPooledConnection";
    private static final String DRUID_POOL_STATEMENT_CLASS = "com/alibaba/druid/pool/DruidPooledStatement";
    private static final String SERVLET_CLASS = "javax/servlet/http/HttpServlet";

    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (loader == null || className == null) {
                    return null;
                }
                byte[] result = null;
                String inputClassName = className.replace("/", ".");
                if (SERVLET_CLASS.equals(className)) {
                    result = ServletAgent.buildAgentClassByte(loader, inputClassName);
                } else if (DRUID_POOL_CONN_CLASS.equals(className)) {
                    result = DruidPoolAgent.buildDruidPoolConnectClassByte(inputClassName);
                } else if (DRUID_POOL_STATEMENT_CLASS.equals(className)) {
                    result = DruidPoolAgent.buildDruidPoolStatementClassByte(inputClassName);
                } else if (MOCK_CLUSTER_INVOKER.equals(className)) {
                    result = DubboAgent.buildDubboAgentClass(inputClassName);
                }
                return result;
            }
        }, true);
    }
}
